// Date:   Mon Aug 11 09:48:46 PM 2025
// Mail:   lunar_ubuntu@qq.com
// Author: https://github.com/xiaoqixian

#include <algorithm>
#include <atomic>
#include <cassert>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <unordered_set>

#include "common.hpp"
#include "config.hpp"
#include "datapipe.hpp"
#include "detector.hpp"
#include "yaml-cpp/node/parse.h"

static bool canceled = false;

struct Stats {
  float mean = 0.0f;
  float stddev = 0.0f;
  size_t min = 0;
  size_t max = 0;
  float median = 0.0f;
  float q1 = 0.0f;
  float q3 = 0.0f;
  float iqr = 0.0f;
  float p90 = 0.0f;
  float p99 = 0.0f;
};

void signal_handler(int) {
  canceled = true;
  exit(0);
}

[[maybe_unused]]
static void save_cv_mat(cv::Mat const& mat, int channel_id) {
  namespace fs = std::filesystem;
  auto folder = cv::format("channel%d", channel_id);
  if (!fs::exists(folder)) {
    fs::create_directory(folder);
  }

  auto filename = cv::format("%s/%s.png", folder.c_str(), timestr().c_str());
  cv::imwrite(filename, mat);
}

static Stats compute_stats(const std::vector<size_t>& data) {
  Stats s;
  if (data.empty()) return s;

  std::vector<size_t> v(data.begin(), data.end());
  std::sort(v.begin(), v.end());

  size_t count = v.size();
  s.min = v.front();
  s.max = v.back();

  // mean & stddev
  double sum = std::accumulate(v.begin(), v.end(), 0.0);
  double sum_sq = 0.0;
  for (auto x : v) sum_sq += static_cast<double>(x) * x;
  s.mean = static_cast<float>(sum / count);
  s.stddev = static_cast<float>(std::sqrt(sum_sq / count - (sum / count)*(sum / count)));

  // median
  if (count % 2 == 1)
    s.median = static_cast<float>(v[count/2]);
  else
    s.median = static_cast<float>((v[count/2 - 1] + v[count/2]) / 2.0);

  // quartiles
  size_t mid = count / 2;
  s.q1 = (mid % 2 == 1) ? static_cast<float>(v[mid/2])
                         : static_cast<float>((v[mid/2 -1] + v[mid/2])/2.0);
  s.q3 = ((count - mid) % 2 == 1) ? static_cast<float>(v[mid + (count-mid)/2])
                                  : static_cast<float>((v[mid + (count-mid)/2 -1] + v[mid + (count-mid)/2])/2.0);
  s.iqr = s.q3 - s.q1;

  // percentiles
  s.p90 = static_cast<float>(v[static_cast<size_t>(0.9*count)]);
  s.p99 = static_cast<float>(v[static_cast<size_t>(0.99*count)]);

  return s;
}

struct ChannelSeqPair {
  int channel_id;
  int emit_seq;

  ChannelSeqPair(int ch, int e) noexcept: channel_id(ch), emit_seq(e) {}

  bool operator==(const ChannelSeqPair& other) const noexcept {
    return channel_id == other.channel_id && emit_seq == other.emit_seq;
  }
};

// Hash specialization
namespace std {
  template <>
  struct hash<ChannelSeqPair> {
    size_t operator()(const ChannelSeqPair& p) const noexcept {
      // A simple but effective hash combine
      size_t h1 = std::hash<int>{}(p.channel_id);
      size_t h2 = std::hash<int>{}(p.emit_seq);
      return h1 ^ (h2 + 0x9e3779b9 + (h1 << 6) + (h1 >> 2));
    }
  };
}

int main(int argc, char** argv) {
  if (argc != 2) {
    printf("Usage: ./build/yolov5_openvino config.json\n");
    exit(1);
  }

  std::string config_path = argv[1];
  Config config = YAML::LoadFile(config_path);

  std::unordered_set<int> detect_labels;
  for (auto const label: config["detect-labels"]) {
    detect_labels.emplace(label.as<int>());
  }

  DataPipe<ActionShot> shot_stream("shot", 1000);
  shot_stream.set_mode(DiscardMode::BLOCK);

  Context ctx;
  ctx.set("shot_stream", &shot_stream);

  Detector detector(config, ctx);

  auto& sampler = detector.sampler_;

  // captured shot classfied by label
  std::unordered_map<int, std::unordered_set<ChannelSeqPair>> shot_cap;
  // captured shot delay classfied by label
  std::unordered_map<int, std::vector<size_t>> delay_by_label;

  detector.start();

  // signal(SIGINT, signal_handler);

  std::atomic_bool stopped = false;

  auto out_th = std::thread([&]() {
    while (!stopped.load(std::memory_order_relaxed)) {
      auto shot_opt = shot_stream.block_pop_front_for(std::chrono::milliseconds(10));
      if (!shot_opt.has_value()) continue;
      auto& shot = shot_opt.value();
      // printf("shot with label/emit_seq/channel = %d/%d/%d\n", shot.label, shot.emit_seq, shot.channel_id);
      shot_cap[shot.label].emplace(shot.channel_id, shot.emit_seq);
      delay_by_label[shot.label].emplace_back(shot.shot_ms - shot.ms_ts);
    }

    std::vector<ActionShot> buffer;
    shot_stream.pop_front_all(buffer);
    for (auto& shot: buffer) {
      // printf("shot with label/emit_seq/channel = %d/%d/%d\n", shot.label, shot.emit_seq, shot.channel_id);
      shot_cap[shot.label].emplace(shot.channel_id, shot.emit_seq);
      delay_by_label[shot.label].emplace_back(shot.shot_ms - shot.ms_ts);
    }
  });

  // detector.stop();
  detector.wait();

  stopped = true;
  if (out_th.joinable()) out_th.join();

  for (auto const node: config["events"]) {
    auto const label = node["label"].as<int>();
    if (detect_labels.count(label) == 0) continue;
    auto const event = node["event"].as<std::string>();

    const size_t total_cap = shot_cap[label].size(), 
          total_emit = sampler.label_emit_cnt_[label].load(std::memory_order_relaxed);
    auto delay_stats = compute_stats(delay_by_label[label]);
    printf("%s: %zd/%zd, ", event.c_str(), total_cap, total_emit);
    printf(
      "delay avg/iqr/p90 = %f/%f/%f\n", 
      static_cast<double>(delay_stats.mean), 
      static_cast<double>(delay_stats.iqr), 
      static_cast<double>(delay_stats.p90)
    );
  }
}
